#!/usr/bin/env python3

import os
from Tools.process_files import process_file
from Tools.build_tree import build_tree
from Tools.analyze_tree import analyze_subroutine_to_string
from Tools.print_tree import generate_relationships, locate_subroutines_in_files
from Tools.analyze_tree import tree_to_string
from Tools.print_info import print_configuration, print_results, set_analysis_parameters
import argparse

# Set debug flag for the entire application
DEBUG = False

def compare_global_variables(root_name, target_subroutine, relationships, relationship_labels=None, 
                            aliases=None, filter_depth=2, filter_mode="down", src_directory=None, 
                            include_steps=False, include_patterns=None, exclude_patterns=None,
                            up_max_search_depth=5, gloabal_vars_label=None, compiler_type="intel"):
    """
    Compare global variables between root function and target subroutine
    
    Returns:
        Tuple containing the comparison analysis and both analysis strings
    """
    lines = []
    lines.append(f"\n===== Global Variable Comparison Analysis: {root_name} vs {target_subroutine} =====")
    
    # Analyze root function and target subroutine
    analysis_root_str, globals_root, *root_rest = analyze_subroutine_to_string(
        root_name, relationships, relationship_labels, aliases, 
        filter_depth=filter_depth, filter_mode=filter_mode, src_directory=src_directory,
        include_steps=include_steps, include_patterns=include_patterns, 
        exclude_patterns=exclude_patterns, root_name=root_name,
        up_max_search_depth=up_max_search_depth, analyze_globals=True,
        gloabal_vars_label=gloabal_vars_label,
        compiler_type=compiler_type
    )
    
    if DEBUG:
        print(f"Debug in compare_global_variables: globals_root = {globals_root}")
    
    analysis_sub_str, globals_subroutine, *sub_rest = analyze_subroutine_to_string(
        target_subroutine, relationships, relationship_labels, aliases, 
        filter_depth=filter_depth, filter_mode=filter_mode, src_directory=src_directory,
        include_steps=include_steps, include_patterns=include_patterns, 
        exclude_patterns=exclude_patterns, root_name=root_name,
        up_max_search_depth=up_max_search_depth, analyze_globals=True,
        gloabal_vars_label=gloabal_vars_label,
        compiler_type=compiler_type
    )
    
    if DEBUG:
        print(f"Debug in compare_global_variables: globals_subroutine = {globals_subroutine}")
    
    # Find variables in root that aren't in the subroutine
    root_only_vars = set(globals_root) - set(globals_subroutine)
    
    lines.append(f"\nVariables written in {root_name} but not written in {target_subroutine}:")
    if root_only_vars:
        for var in sorted(root_only_vars):
            lines.append(f"  - {var}")
    else:
        lines.append("  No additional variables found")
    
    return ("\n".join(lines), analysis_root_str, analysis_sub_str)




def main(src_directory=None, isRecursive=False):
    """Main function, processes F90 files and returns the relationship graph and program name list"""
    if src_directory is None:
        src_directory = ['.']
    
    # Collect all F90 files
    f90_files = []
    for directory in src_directory:
        if isRecursive:
            for root, _, files in os.walk(directory):
                for file in files:
                    if file.lower().endswith(('.f90', '.f')) or file.endswith(('.F90', '.F')):
                        f90_files.append(os.path.join(root, file))
        else:
            f90_files.extend([os.path.join(directory, f) for f in os.listdir(directory) if f.lower().endswith(('.f90', '.f')) or f.endswith(('.F90', '.F'))])
    
    # Process all files and collect relationships
    all_relationships = {}
    all_chunks = []
    all_aliases = {}
 
    for file in f90_files:
        chunks, file_aliases = process_file(file)
        all_chunks.extend(chunks)
        all_aliases.update(file_aliases)
        
    # Process relationships after collecting all aliases
    relationships = generate_relationships(all_chunks, all_aliases)
    all_relationships.update(relationships)

    # Collect all program names, not just the first one
    root_names = [chunk['name'] for chunk in all_chunks if chunk['type'] == 'program']
    print(f"Detected program name: {root_names}")
    
    # If no programs are found, root_names will be an empty list
    return all_relationships, root_names, all_aliases


if __name__ == "__main__":
    # Analysis flags
    analysis_flag = "none"  # Can be set to "drp4dvar" or "gsi"
    filter_depth = {"down": 5, "up": 1}  # Set different depths for different filter modes
    filter_mode = "none"  # Filter mode, 'down', 'up' or 'both'
    filter_mode_sub = "none"
    #CODE_DIR = '../../var4_test/'
    
    # Use command line argument parsing
    src_default = ['.']
    target_default = None
    parser = argparse.ArgumentParser(description="Fortran Flow Map: Analyze Fortran code structure and call relationships")
    parser.add_argument('--src', nargs='+', default=src_default, help='Source directories to scan for Fortran files')
    parser.add_argument('--target', default=target_default, help='Target subroutine name for specific analysis')
    parser.add_argument('--root', default=None, help='Root program name for analysis')
    parser.add_argument('--debug', action='store_true', help='Enable debug output')
    parser.add_argument('--filter', default=None, help='Comma-separated list of additional function names to filter out (e.g., "grapes_error_fatal,global2patch")')
    parser.add_argument('--output_dir', default='Output', help='Custom output directory for generated files')
    parser.add_argument('--all', action='store_true', help='Process all available root program names')
    args = parser.parse_args()
    
    # Set debug flag from command line argument
    DEBUG = args.debug
    
    # Use command line arguments if provided
    src_directory = args.src
    target_subroutine = args.target
    user_root_name = args.root
    output_dir = args.output_dir
    process_all = args.all
    
    print_configuration(analysis_flag, filter_depth, filter_mode, src_directory, user_root_name, target_subroutine)
    
    result_label = None
    result, root_names, aliases = main(src_directory=src_directory, isRecursive=True)
    # Continue with user-specified root program name or process all
    from Tools.check_root_name import check_root_name
    from Tools.filter_def_names import func_names
    
    gloabal_vars = None
    compiler_type = "intel"
    RAG = False
    
    # Process additional functions to filter if specified via --filter
    if args.filter:
        additional_filters = [func.strip() for func in args.filter.split(',')]
        for func in additional_filters:
            if func and func not in func_names:
                func_names.append(func)
                if DEBUG:
                    print(f"Added {func} to filter list")
    
    if process_all:
        print(f"Processing all detected root programs: {root_names}")
        for root in root_names:
            root_name = check_root_name(root, result, root_names)
            # Call compare_global_variables to get both analyses and comparison
            comparison_analysis, root_analysis, target_analysis = compare_global_variables(
                root_name, target_subroutine if target_subroutine else root_name, result, result_label, aliases, 
                filter_depth=filter_depth, filter_mode=filter_mode_sub, 
                src_directory=src_directory, include_steps=True, 
                exclude_patterns=func_names, up_max_search_depth=4,
                gloabal_vars_label=gloabal_vars,
                compiler_type=compiler_type
            )
            
            # Debug information
            if DEBUG:
                print(f"Debug: Root name: {root_name}")
                print(f"Debug: Result keys: {list(result.keys())}")
                print(f"Debug: Number of functions: {len(result)}")
                
                # Print detailed information about the root entry
                if root_name in result:
                    print(f"Debug: {root_name} calls in result dictionary: {len(result[root_name])}")
                    for i, call in enumerate(result[root_name]):
                        print(f"Debug: Call {i+1}: {call['func']} with args: {call.get('args', '')}")
                else:
                    print(f"Debug: {root_name} not found in result dictionary")
                
                # Import locate_subroutines_in_files function to check subroutine locations
                subroutine_locations = locate_subroutines_in_files(src_directory, recursive=True)
                print(f"Debug: Found location for {root_name}: {root_name in subroutine_locations}")
                if root_name in subroutine_locations:
                    print(f"Debug: Location: {subroutine_locations[root_name]}")
            
            # Generate and save tree visualization
            tree_file = f"{output_dir}/tree_{root_name}.txt"
            try:
                # Ensure the directory exists before writing the file
                os.makedirs(os.path.dirname(tree_file), exist_ok=True)
                if root_name in result:
                    # Use "none" as filter_mode to show all function calls
                    tree_main = build_tree(root_name, result, result_label, aliases, 
                                     filter_depth=filter_depth, filter_mode="none", 
                                     exclude_patterns=func_names)
                    tree_str, _ = tree_to_string(tree_main, include_steps=False)
                    with open(tree_file, 'w', encoding='utf-8') as f:
                        f.write("\n".join([str(line) for line in tree_str]))
                    print(f"Successfully wrote tree visualization to {tree_file}")
                else:
                    print(f"Error: Could not generate tree file - {root_name} not found in results")
            except Exception as e:
                print(f"Error generating tree file for {root_name}: {str(e)}")
                
            
            if RAG:
                # Write both analyses to separate files
                analysis_root_file = f"{output_dir}/analysis_{root_name}.txt"
                # Ensure the directory exists before writing the file
                os.makedirs(os.path.dirname(analysis_root_file), exist_ok=True)
                with open(analysis_root_file, 'w', encoding='utf-8') as f:
                    f.write(root_analysis)
                print(f"Successfully wrote root analysis to {analysis_root_file}")
                
                target_name = target_subroutine if target_subroutine else root_name
                analysis_target_file = f"{output_dir}/analysis_{target_name}.txt"
                # Ensure the directory exists before writing the file
                os.makedirs(os.path.dirname(analysis_target_file), exist_ok=True)
                with open(analysis_target_file, 'w', encoding='utf-8') as f:
                    f.write(target_analysis)
                print(f"Successfully wrote target analysis to {analysis_target_file}")
                
                # Write comparison to file
                comparison_file = f"{output_dir}/comparison_{root_name}_vs_{target_name}.txt"
                # Ensure the directory exists before writing the file
                os.makedirs(os.path.dirname(comparison_file), exist_ok=True)
                with open(comparison_file, 'w', encoding='utf-8') as f:
                    f.write(comparison_analysis)
                print(f"Successfully wrote comparison analysis to {comparison_file}")
    else:
        if target_subroutine is None:
            target_subroutine = user_root_name
            print(f"The target subroutine name is not specified, and the root program name is used by default: {user_root_name}")
        # Convert names to lowercase for case-insensitive comparison
        root_names_lower = [name.lower() for name in root_names]
        user_root_name_lower = user_root_name.lower() if user_root_name else None
        # Adjust user_root_name to match the case of the detected program
        if user_root_name_lower in root_names_lower:
            index = root_names_lower.index(user_root_name_lower)
            user_root_name = root_names[index]
            print(f"Adjust the case of the root program name to match the name in the code: {user_root_name}")
        root_name = check_root_name(user_root_name, result, root_names)

        # Call compare_global_variables to get both analyses and comparison
        comparison_analysis, root_analysis, target_analysis = compare_global_variables(
            root_name, target_subroutine, result, result_label, aliases, 
            filter_depth=filter_depth, filter_mode=filter_mode_sub, 
            src_directory=src_directory, include_steps=True, 
            exclude_patterns=func_names, up_max_search_depth=4,
            gloabal_vars_label=gloabal_vars,
            compiler_type=compiler_type
        )
        
        # Debug information
        if DEBUG:
            print(f"Debug: Root name: {root_name}")
            print(f"Debug: Result keys: {list(result.keys())}")
            print(f"Debug: Number of functions: {len(result)}")
            
            # Print detailed information about the grapes entry
            if root_name in result:
                print(f"Debug: {root_name} calls in result dictionary: {len(result[root_name])}")
                for i, call in enumerate(result[root_name]):
                    print(f"Debug: Call {i+1}: {call['func']} with args: {call.get('args', '')}")
            else:
                print(f"Debug: {root_name} not found in result dictionary")
            
            # Import locate_subroutines_in_files function to check subroutine locations
            subroutine_locations = locate_subroutines_in_files(src_directory, recursive=True)
            print(f"Debug: Found location for grapes: {root_name in subroutine_locations}")
            if root_name in subroutine_locations:
                print(f"Debug: Location: {subroutine_locations[root_name]}")
        
        # Generate and save tree visualization
        tree_file = f"{output_dir}/tree_{root_name}.txt"
        try:
            # Ensure the directory exists before writing the file
            os.makedirs(os.path.dirname(tree_file), exist_ok=True)
            if root_name in result:
                # Use "none" as filter_mode to show all function calls
                tree_main = build_tree(root_name, result, result_label, aliases, 
                                 filter_depth=filter_depth, filter_mode="none", 
                                 exclude_patterns=func_names)
                tree_str, _ = tree_to_string(tree_main, include_steps=False)
                with open(tree_file, 'w', encoding='utf-8') as f:
                    f.write("\n".join([str(line) for line in tree_str]))
                print(f"Successfully wrote tree visualization to {tree_file}")
            else:
                print(f"Error: Could not generate tree file - {root_name} not found in results")
        except Exception as e:
            print(f"Error generating tree file: {str(e)}")
        
        if RAG:
            # Write both analyses to separate files
            analysis_root_file = f"{output_dir}/analysis_{root_name}.txt"
            # Ensure the directory exists before writing the file
            os.makedirs(os.path.dirname(analysis_root_file), exist_ok=True)
            with open(analysis_root_file, 'w', encoding='utf-8') as f:
                f.write(root_analysis)
            print(f"Successfully wrote root analysis to {analysis_root_file}")
            
            analysis_target_file = f"{output_dir}/analysis_{target_subroutine}.txt"
            # Ensure the directory exists before writing the file
            os.makedirs(os.path.dirname(analysis_target_file), exist_ok=True)
            with open(analysis_target_file, 'w', encoding='utf-8') as f:
                f.write(target_analysis)
            print(f"Successfully wrote target analysis to {analysis_target_file}")
            
            # Write comparison to file
            comparison_file = f"{output_dir}/comparison_{root_name}_vs_{target_subroutine}.txt"
            # Ensure the directory exists before writing the file
            os.makedirs(os.path.dirname(comparison_file), exist_ok=True)
            with open(comparison_file, 'w', encoding='utf-8') as f:
                f.write(comparison_analysis)
            print(f"Successfully wrote comparison analysis to {comparison_file}")
